"""
孙守宇(2769919224) 2021/9/6 20:10:29
月考结束，不论升班还是末班，同学学习人工智能“永远在路上”

今晚自习，大家研究下，这个有关“验证码”项目的链接，项目讲述比较全面
https://www.dtmao.cc/news_show_4159411.shtml
深度学习100例-卷积神经网络（CNN）识别验证码

"""

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, optimizers, losses
from sklearn.model_selection import train_test_split
import os

np.random.seed(3)
tf.random.set_seed(3)
FILE_NAME = os.path.basename(__file__)

vcode_dir_train = '../../../../large_data/DL1/_many_files/vcode_data/train'
vcode_dir_test = '../../../../large_data/DL1/_many_files/vcode_data/test'

def load_dir(vcode_dir):
    names = os.listdir(vcode_dir)
    n_names = len(names)
    names = np.array(names)
    idx = np.random.permutation(n_names)  # for shuffle
    names = names[idx]
    x = []
    y = []
    for name in names:
        digits, ext = os.path.splitext(name)
        if ext.lower() != '.jpg':
            continue
        digits_list = list(digits)
        if len(digits_list) != 4:
            continue
        digits_arr = []
        try:
            for di in digits_list:
                digits_arr.append(int(di))
        except ValueError as ex:
            continue

        y.append(digits_arr)

        path = os.path.join(vcode_dir, name)
        img = cv.imread(path, cv.IMREAD_GRAYSCALE)
        x.append(img)

    x = np.float32(x) / 255.
    y = np.uint8(y)
    x = np.expand_dims(x, 3)
    y = keras.utils.to_categorical(y, 10)
    return x, y

print('Loading pictures ...')
x_train, y_train = load_dir(vcode_dir_train)
x_test, y_test = load_dir(vcode_dir_test)
print('Loaded.')

print('x_train:', x_train.shape)
print('y_train:', y_train.shape)
print('x_test:', x_test.shape)
print('y_test:', y_test.shape)

model = keras.Sequential([
    layers.Conv2D(6, (5, 5)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(16, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(32, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(64, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Flatten(),

    # layers.Dropout(0.1),
    layers.Dense(120, activation=activations.relu),
    # layers.Dropout(0.1),  # ATTENTION Drop out.
    layers.Dense(84, activation=activations.relu),
    # layers.Dropout(0.1),
    layers.Dense(40, activation=None),

    layers.Reshape((4, 10)),
    layers.Softmax(),
])
model.build(input_shape=(None, 60, 160, 1))
model.summary(line_length=120)


def my_acc(y_true, y_pred):
    y_true = tf.argmax(y_true, axis=2)
    y_pred = tf.argmax(y_pred, axis=2)
    eq = y_true == y_pred
    eq = tf.reduce_all(eq, axis=1)
    acc = tf.reduce_mean(tf.cast(eq, dtype=tf.float32))
    return acc


def my_acc_digit(y_true, y_pred):
    y_true = tf.argmax(y_true, axis=2)
    y_pred = tf.argmax(y_pred, axis=2)
    eq = y_true == y_pred
    acc = tf.reduce_mean(tf.cast(eq, dtype=tf.float32))
    return acc


model.compile(loss='categorical_crossentropy',
              optimizer=optimizers.Adam(lr=0.001),
              metrics=[my_acc, my_acc_digit, 'accuracy'])

# v2.1 accuracy 0.8578125238418579
VER = 'v5.22'
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
os.makedirs(SAVE_DIR, exist_ok=True)
SAVE_PREFIX = SAVE_DIR + '/weights'
LOG_DIR = os.path.join('_log', FILE_NAME, VER)

if len(os.listdir(SAVE_DIR)) > 0:
    model.load_weights(SAVE_PREFIX)
else:

    history=model.fit(
        x_train, y_train, batch_size=64, epochs=25, validation_split=0.1,
        callbacks=[keras.callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0)],
    )
    model.save_weights(SAVE_PREFIX)

score=model.evaluate(x_test, y_test, batch_size=64)
print(score)
print('loss',score[0])
print('my_acc',score[1])
print('my_acc_digit',score[2])
print('accuracy',score[3])

pred = model.predict(x_test)


def target2label(vec):
    return ''.join([str(i) for i in vec])


offset = 0
spr = 5
spc = 5
spn = 0
plt.figure(figsize=(14, 7))
for i in range(spr * spc):
    idx = offset + i
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.axis('off')
    img = x_test[idx]
    img = np.squeeze(img)
    plt.imshow(img, cmap='gray')
    y_true = y_test[idx].argmax(axis=1)
    y_pred = pred[idx].argmax(axis=1)
    y_true = target2label(y_true)
    y_pred = target2label(y_pred)
    marker = 'V' if y_true == y_pred else 'X'
    title = target2label(f'{y_true}: {y_pred} ({marker})')
    plt.title(title)

plt.show()
